# A combination of geom_line and geom_ribbon  with sensible defaults for displaying multiple bands
#
# Author: mjskay
###############################################################################


# Names that should be suppressed from global variable check by codetools
# Names used broadly should be put in _global_variables.R
globalVariables(c(".lower", ".upper", ".width"))


#' Line + multiple uncertainty ribbon plots (ggplot geom)
#'
#' A combination of [geom_line()] and [geom_ribbon()] with default aesthetics
#' designed for use with output from [point_interval()].
#'
#' `geom_lineribbon` is a combination version of a [geom_line()], and `geom_ribbon` designed for use
#' with output from [point_interval()]. This geom sets some default aesthetics equal to the `.width`
#' column generated by the `point_interval` family of functions, making them
#' often more convenient than a vanilla [geom_ribbon()] + [geom_line()].
#'
#' Specifically, `geom_lineribbon` acts as if its default aesthetics are
#' `aes(fill = forcats::fct_rev(ordered(.width)))`.
#'
#' @inheritParams ggplot2::geom_line
#' @param ...  Other arguments passed to [layer()].
#' @param step Should the line/ribbon be drawn as a step function? One of: `FALSE` (do not draw as a step
#' function, the default), `TRUE` (draw a step function using the `"mid"` approach), `"mid"`
#' (draw steps midway between adjacent x values), `"hv"` (draw horizontal-then-vertical steps), `"vh"`
#' (draw as vertical-then-horizontal steps). `TRUE` is an alias for `"mid"` because for a step function with
#' ribbons, `"mid"` is probably what you want (for the other two step approaches the ribbons at either the
#' vert first or vert last x value will not be visible).
#' @param orientation Whether this geom is drawn horizontally (`"horizontal"`) or
#' vertically (`"vertical"`). The default, `NA`, automatically detects the orientation based on how the
#' aesthetics are assigned, and should generally do an okay job at this. When horizontal (resp. vertical),
#' the geom uses the `y` (resp. `x`) aesthetic to identify different groups, then for each group uses
#' the `x` (resp. `y`) aesthetic and the `thickness` aesthetic to draw a function as an slab, and draws
#' points and intervals horizontally (resp. vertically) using the `xmin`, `x`, and `xmax` (resp.
#' `ymin`, `y`, and `ymax`) aesthetics. For compatibility with the base
#' ggplot naming scheme for `orientation`, `"x"` can be used as an alias for `"vertical"` and `"y"` as an alias for
#' `"horizontal"` (tidybayes had an `orientation` parameter before ggplot did, and I think the tidybayes naming
#' scheme is more intuitive: `"x"` and `"y"` are not orientations and their mapping to orientations is, in my
#' opinion, backwards; but the base ggplot naming scheme is allowed for compatibility).
#' @return A [ggplot2::Geom] representing a combined line+uncertainty ribbon geometry which can
#' be added to a [ggplot()] object.
#' @author Matthew Kay
#' @seealso See [stat_lineribbon()] for a version that does summarizing of samples into points and intervals
#' within ggplot. See [geom_pointinterval()] for a similar geom intended
#' for point summaries and intervals. See [geom_ribbon()] and [geom_line()] for the geoms this is
#' based on.
#' @examples
#'
#' library(dplyr)
#' library(ggplot2)
#'
#' theme_set(theme_ggdist())
#'
#' tibble(x = 1:10) %>%
#'   group_by_all() %>%
#'   do(tibble(y = rnorm(100, .$x))) %>%
#'   median_qi(.width = c(.5, .8, .95)) %>%
#'   ggplot(aes(x = x, y = y, ymin = .lower, ymax = .upper)) +
#'   # automatically uses aes(fill = fct_rev(ordered(.width)))
#'   geom_lineribbon() +
#'   scale_fill_brewer()
#'
#' @importFrom forcats fct_rev
#' @import ggplot2
#' @export
geom_lineribbon = function(
  mapping = NULL,
  data = NULL,
  stat = "identity",
  position = "identity",
  ...,

  step = FALSE,
  orientation = NA,

  na.rm = FALSE,
  show.legend = NA,
  inherit.aes = TRUE
) {

  l = layer(
    data = data,
    mapping = mapping,
    stat = stat,
    geom = GeomLineribbon,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      step = step,
      orientation = orientation,
      na.rm = na.rm,
      ...
    )
  )

  add_default_computed_aesthetics(l,
    aes(fill = forcats::fct_rev(ordered(.width)))
  )
}

draw_key_lineribbon = function(data, params, size) {
  if (!is.null(data$fill_ramp) && is.na(data$fill)) {
    data$fill = "gray65"
  }
  data$fill = apply_colour_ramp(data$fill, data$fill_ramp)

  if (is.na(data$fill)) {
    draw_key_path(data, params, size)
  } else {
    draw_key_rect(data, params, size)
  }
}

#' @rdname ggdist-ggproto
#' @format NULL
#' @usage NULL
#' @importFrom purrr map map_dbl reduce
#' @import ggplot2
#' @export
GeomLineribbon = ggproto("GeomLineribbon", Geom,
  default_aes = aes(
    colour = "black",
    size = 1.25,
    linetype = 1,
    shape = 19,
    fill = NA,
    fill_ramp = NULL,
    alpha = NA,
    stroke = 1
  ),

  draw_key = draw_key_lineribbon,

  required_aes = c("x", "y"),

  optional_aes = c("ymin", "ymax", "xmin", "xmax", "fill_ramp"),

  extra_params = c("step", "orientation", "na.rm"),

  default_params = list(
    step = FALSE,
    orientation = NA,
    na.rm = FALSE
  ),

  setup_params = function(self, data, params) {
    params = defaults(params, self$default_params)

    # detect orientation
    params$flipped_aes = get_flipped_aes(data, params,
      range_is_orthogonal = TRUE, ambiguous = TRUE, group_has_equal = TRUE
    )
    params$orientation = get_orientation(params$flipped_aes)

    params
  },

  setup_data = function(self, data, params) {
    #set up orientation
    data$flipped_aes = params$flipped_aes

    data
  },

  draw_panel = function(self, data, panel_scales, coord,
    step = self$default_params$step,
    orientation = self$default_params$orientation,
    flipped_aes = FALSE
  ) {
    define_orientation_variables(orientation)

    data$fill = apply_colour_ramp(data$fill, data$fill_ramp)

    # ribbons do not autogroup by color/fill/linetype, so if someone groups by changing the color
    # of the line or by setting fill, the ribbons might give an error. So we will do the
    # grouping ourselves
    grouping_columns = names(data) %>%
      intersect(c("colour", "fill", "linetype", "group"))

    # draw as a step function if requested
    if (step == TRUE) step = "mid"
    if (step != FALSE) data = ddply_(data, grouping_columns, stepify, x = y, direction = step)

    # draw all the ribbons
    ribbon_grobs = data %>%
      dlply_(grouping_columns, function(d) {
        group_grobs = list(GeomRibbon$draw_panel(transform(d, size = NA), panel_scales, coord, flipped_aes = flipped_aes))
        list(
          width = mean(abs(d[[xmax]] - d[[xmin]])),
          grobs = group_grobs
        )
      })

    # this is a slightly hackish approach to getting the draw order correct for the common
    # use case of fit lines / curves: draw the ribbons in order from largest mean width to
    # smallest mean width, so that the widest intervals are on the bottom.
    ribbon_grobs = ribbon_grobs[order(-map_dbl(ribbon_grobs, "width"))] %>%
      map("grobs") %>%
      reduce(c)

    # now draw all the lines
    line_grobs = data %>%
      dlply_(grouping_columns, function(d) {
        if (!is.null(d[[x]])) {
          list(GeomLine$draw_panel(d, panel_scales, coord))
        } else {
          list()
        }
      })

    line_grobs = reduce(line_grobs, c)

    grobs = c(ribbon_grobs, line_grobs)

    ggname("geom_lineribbon",
      gTree(children = do.call(gList, grobs))
    )
  }
)


# helpers -----------------------------------------------------------------

#' @importFrom dplyr lag lead
stepify = function(df, x = "x", direction = "hv") {
  n = nrow(df)

  # sort by x and double up all rows in the data frame
  step_df = df[rep(order(df[[x]]), each = 2),]

  if (direction == "hv") {
    # horizontal-to-vertical step => lead x and drop last row
    step_df[[x]] = lead(step_df[[x]])
    step_df[-2*n,]
  } else if (direction == "vh") {
    # vertical-to-horizontal step => lag x and drop first row
    step_df[[x]] = lag(step_df[[x]])
    step_df[-1,]
  } else if (direction == "mid") {
    # mid step => last value in each pair is matched with the first value in the next pair,
    # then we set their x position to their average.
    # Need to repeat the last value one more time to make it work
    step_df[2*n + 1,] = step_df[2*n,]

    x_i = seq_len(n)*2
    mid_x = (step_df[x_i, x] + step_df[x_i + 1, x]) / 2

    step_df[x_i, x] = mid_x
    step_df[x_i + 1, x] = mid_x
    step_df
  }
}
